# 探究植物生长受光照、水分和温度三种因素的影响情况，光照、水分、温度都分为低和高两种水平
# 植物的生长情况由植物生长高度和叶子数量来度量
import multiprocessing

import math
import pandas as pd
from cadCAD.configuration.utils import config_sim
from cadCAD import configs, Experiment
from cadCAD.engine import ExecutionContext, Executor, ExecutionMode
import matplotlib.pyplot as plt
import random
from pprint import pprint

init_state = {
    'height': 0,
    'leaf_num': 0
}
params = {
    'N': 1,
    'T': 10,
    'sunlight': ['low', 'high'],
    'water': ['low', 'high'],
    'temperature': ['low', 'high']
}


def dfs(x, param, index, capacity):
    if x == 3:
        s = 0
        for key in capacity.keys():
            capacity[key].append(param[key][index[s]])
            s = s + 1
        return capacity
    if x == 0:
        for i in range(len(param['sunlight'])):
            index[x] = i
            capacity = dfs(x + 1, param, index, capacity)
    if x == 1:
        for i in range(len(param['water'])):
            index[x] = i
            capacity = dfs(x + 1, param, index, capacity)
    if x == 2:
        for i in range(len(param['temperature'])):
            index[x] = i
            capacity = dfs(x + 1, param, index, capacity)
    return capacity


def getCapacity(param):
    capacity = {
        'sunlight': [],
        'water': [],
        'temperature': []
    }
    index = []
    for i in range(len(capacity)):
        index.append(0)
    capacity = dfs(0, param, index, capacity)
    return capacity


def add_height(params, substep, state_history, previous_state):
    sunlight = 0.5 if params['sunlight'] == 'low' else 1
    water = 0.5 if params['water'] == 'low' else 1
    temperature = 0.5 if params['temperature'] == 'low' else 1

    value = sunlight * water * temperature
    return {'height_add_value': value}


def add_leaf(params, substep, state_history, previous_state):
    sunlight = 0.5 if params['sunlight'] == 'low' else 1
    water = 0.5 if params['water'] == 'low' else 1
    temperature = 0.5 if params['temperature'] == 'low' else 1

    value = sunlight + water + temperature
    return {'leaf_add_value': value}


def update_height(params, substep, state_history, previous_state, policy_input):
    add_value = policy_input['height_add_value']
    height = previous_state['height']
    new_height = height + add_value
    # print(new_height)
    return 'height', new_height


def update_leaf_num(params, substep, state_history, previous_state, policy_input):
    add_value = policy_input['leaf_add_value']
    leaf_num = previous_state['leaf_num']
    new_leaf_num = leaf_num + add_value
    return 'leaf_num', new_leaf_num


psubs = [
    {
        'policies': {
            'height_add_value': add_height,
            'leaf_add_value': add_leaf
        },
        'variables': {
            'height': update_height,
            'leaf_num': update_leaf_num,
        }
    }
]


def run():
    del configs[:]

    capacity = getCapacity(params)
    print('全因子设计表')
    pprint(capacity)
    sim_config = config_sim({
        'N': params['N'],
        'T': range(params['T']),
        'M': capacity
    })
    exp = Experiment()
    exp.append_configs(
        sim_configs=sim_config,
        initial_state=init_state,
        partial_state_update_blocks=psubs,
    )

    exec_mode = ExecutionMode()
    exec_context = ExecutionContext(exec_mode.multi_mode)  # 不写默认为local_mode
    run = Executor(exec_context=exec_context, configs=exp.configs)

    (raw_result, tensor, sessions) = run.execute()
    df = pd.DataFrame(raw_result)
    df.set_index(['simulation', 'run', 'timestep', 'substep'])
    pprint(df)

    for subset in df.subset.unique():
        df[(df['subset'] == subset)].plot(
            x='timestep', y=['leaf_num', 'height'], marker='o',
            markersize=12,
            markeredgewidth=4, alpha=0.2,
            markerfacecolor='black',
            linewidth=5, figsize=(12, 8),
            title="Factors affecting plant growth",
            ylabel='Number', grid=True,
            fillstyle='none',
            xticks=list(df['timestep'].drop_duplicates()),
            legend=None,
            yticks=list(range(1 + 30))
        )
        plt.show()


if __name__ == "__main__":
    multiprocessing.freeze_support()
    print('全因子设计')
    run()
